import os
import numpy as np
import sys
import time
import datetime
import csv
import shutil

# 添加必要的路径
sys.path.append("/home/konghaomin/TrackEval/scripts")
import run_mot_challenge_func

tracker_name = {
    "bytetrack": "ByteTrack",
    "boosttrack": "BoostTrack",
    "imprassoctrack": "ImprAssocTrack",
    "strongsort": "StrongSORT",
    "ocsort": "OCSort",
    "deepocsort": "DeepOCSort",
    "botsort": "BoTSORT",
    # 可根据需要继续扩展
}


def convert_result_to_str(result):
    """将结果转换为百分比字符串，保留3位小数"""
    return str(round(result * 100, 3))


def evaluate_tracker(tracker_path, tracker_key):
    """评估单个跟踪器并返回其性能指标"""
    display_name = tracker_name.get(tracker_key.lower(), tracker_key)
    print(f"\n正在评估 {display_name}...")

    try:
        res_eval = run_mot_challenge_func.main(
            SPLIT_TO_EVAL="val",
            METRICS=["HOTA", "CLEAR", "Identity"],
            GT_FOLDER="/home/konghaomin/Datasets/MaritimeTrack_20250322_D/DanceTrack/val",
            SEQMAP_FILE="/home/konghaomin/Datasets/MaritimeTrack_20250322_D/DanceTrack/val_seqmap.txt",
            SKIP_SPLIT_FOL=True,
            TRACKERS_TO_EVAL=[""],
            TRACKER_SUB_FOLDER="",
            USE_PARALLEL=True,
            NUM_PARALLEL_CORES=8,
            PLOT_CURVES=False,
            TRACKERS_FOLDER=tracker_path,
        )

        # 提取各种指标
        hota = float(
            np.mean(
                res_eval[0]["MotChallenge2DBox"][""]["COMBINED_SEQ"]["pedestrian"][
                    "HOTA"
                ]["HOTA"]
            )
        )

        deta = float(
            np.mean(
                res_eval[0]["MotChallenge2DBox"][""]["COMBINED_SEQ"]["pedestrian"][
                    "HOTA"
                ]["DetA"]
            )
        )

        assa = float(
            np.mean(
                res_eval[0]["MotChallenge2DBox"][""]["COMBINED_SEQ"]["pedestrian"][
                    "HOTA"
                ]["AssA"]
            )
        )

        mota = res_eval[0]["MotChallenge2DBox"][""]["COMBINED_SEQ"]["pedestrian"][
            "CLEAR"
        ]["MOTA"]

        idf1 = res_eval[0]["MotChallenge2DBox"][""]["COMBINED_SEQ"]["pedestrian"][
            "Identity"
        ]["IDF1"]

        # 转换结果为字符串格式
        hota_str = convert_result_to_str(hota)
        deta_str = convert_result_to_str(deta)
        assa_str = convert_result_to_str(assa)
        mota_str = convert_result_to_str(mota)
        idf1_str = convert_result_to_str(idf1)

        return {
            "tracker": display_name,
            "hota": hota_str,
            "deta": deta_str,
            "assa": assa_str,
            "mota": mota_str,
            "idf1": idf1_str,
            "success": True,
        }
    except Exception as e:
        print(f"评估 {display_name} 时出错: {str(e)}")
        return {"tracker": display_name, "success": False, "error": str(e)}


def main():
    # 结果目录
    result_dir = "/home/konghaomin/boxmot-mt/tracker_output/All_DDETR"

    output_dir_base = "./eval_result"

    # 创建基于目录名和时间的输出目录
    dir_basename = os.path.basename(result_dir)
    current_time = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    output_dir_name = f"{dir_basename}_{current_time}"
    output_dir = os.path.join(output_dir_base, output_dir_name)

    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    print(f"所有输出将保存到: {output_dir}")

    # 获取所有跟踪器目录
    tracker_list = os.listdir(result_dir)
    tracker_dir = [os.path.join(result_dir, tracker) for tracker in tracker_list]

    print(f"找到以下 {len(tracker_list)} 个跟踪器:")
    for i, tracker in enumerate(tracker_list, 1):
        print(f"{i}. {tracker}")

    # 存储所有评估结果
    all_results = []
    success_count = 0

    # 对每个跟踪器执行评估
    for i, (tracker, tracker_path) in enumerate(zip(tracker_list, tracker_dir), 1):
        print(f"\n[{i}/{len(tracker_list)}] 开始评估: {tracker}")

        # 检查是否有data子目录
        if os.path.exists(os.path.join(tracker_path, "data")):
            tracker_path = os.path.join(tracker_path, "data")

        # 执行评估
        result = evaluate_tracker(tracker_path, tracker)

        # 保存结果
        if result["success"]:
            all_results.append(result)
            success_count += 1

            # 将结果写入单独的日志文件（在新目录中）
            with open(
                os.path.join(output_dir, f"Result_{result['tracker']}.log"), "w"
            ) as f:
                f.write(f"HOTA: {result['hota']}\n")
                f.write(f"DetA: {result['deta']}\n")
                f.write(f"AssA: {result['assa']}\n")
                f.write(f"MOTA: {result['mota']}\n")
                f.write(f"IDF1: {result['idf1']}\n")

            print(f"✓ {result['tracker']} 评估完成")
            print(
                f"HOTA: {result['hota']}, DetA: {result['deta']}, AssA: {result['assa']}, MOTA: {result['mota']}, IDF1: {result['idf1']}"
            )
        else:
            print(f"✗ {tracker} 评估失败: {result.get('error', '未知错误')}")

        time.sleep(1)  # 稍微暂停以便查看输出

    # 输出汇总结果
    print(f"\n\n评估完成! 共评估 {len(tracker_list)} 个跟踪器, 成功 {success_count} 个")

    if all_results:
        # 按HOTA分数排序（升序）
        sorted_results = sorted(
            all_results, key=lambda x: float(x["hota"])  # 移除了 reverse=True 参数
        )

        print("\n所有跟踪器性能汇总 (按HOTA升序排列):")
        print(
            f"{'Tracker':<15} {'HOTA':<8} {'DetA':<8} {'AssA':<8} {'MOTA':<8} {'IDF1':<8}"
        )
        print("-" * 60)

        for r in sorted_results:
            print(
                f"{r['tracker']:<15} {r['hota']:<8} {r['deta']:<8} {r['assa']:<8} {r['mota']:<8} {r['idf1']:<8}"
            )

        # 将完整排名写入文件（在新目录中）
        summary_txt_path = os.path.join(output_dir, "tracker_evaluation_summary.txt")
        with open(summary_txt_path, "w") as f:
            f.write(
                f"{'Tracker':<15} {'HOTA':<8} {'DetA':<8} {'AssA':<8} {'MOTA':<8} {'IDF1':<8}\n"
            )
            f.write("-" * 60 + "\n")
            for r in sorted_results:
                f.write(
                    f"{r['tracker']:<15} {r['hota']:<8} {r['deta']:<8} {r['assa']:<8} {r['mota']:<8} {r['idf1']:<8}\n"
                )

        # 额外保存为CSV格式
        csv_path = os.path.join(output_dir, "tracker_evaluation_summary.csv")
        with open(csv_path, "w", newline="") as csvfile:
            fieldnames = ["Tracker", "HOTA", "DetA", "AssA", "MOTA", "IDF1"]
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

            writer.writeheader()
            for r in sorted_results:
                writer.writerow(
                    {
                        "Tracker": r["tracker"],
                        "HOTA": r["hota"],
                        "DetA": r["deta"],
                        "AssA": r["assa"],
                        "MOTA": r["mota"],
                        "IDF1": r["idf1"],
                    }
                )

        print(f"评估结果摘要已保存到: {summary_txt_path}")
        print(f"CSV格式结果已保存到: {csv_path}")


if __name__ == "__main__":
    main()
